//以当前位置做力控
#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp> // 确保包含 WrenchStamped
#include <tf2_ros/transform_listener.hpp>
#include <tf2_ros/buffer.hpp>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp> // 推荐加上这个以防万一
#include <eigen3/Eigen/Dense>

using namespace std::chrono_literals;

class ForceServoControlNode : public rclcpp::Node
{
public:
    ForceServoControlNode() : Node("force_servo_control_node"), count_(0)
    {
        // 参数初始化
        M_position_ = Eigen::Matrix3d::Identity() * 1.0;  // 质量 M
        K_position_ = Eigen::Matrix3d::Identity() * 18.0; // 刚度 K
        B_position_ = Eigen::Matrix3d::Identity() * 13.0; // 阻尼 B
        force_torque_reference_ = Eigen::Vector3d(0.0, 0.0, 0.0);
        velocity_ = Eigen::Vector3d::Zero();
        dt_ = 0.04; // 与定时器周期一致 (40ms = 0.04s)
        current_force_ = Eigen::Vector3d::Zero();

        // 1. 创建力传感器数据订阅者
        force_sub_ = create_subscription<geometry_msgs::msg::WrenchStamped>(
            "/force_data", 10, std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

        // 2. 初始化TF监听器
        tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
        tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

        // 3. 设置发布者 [关键修改]
        // 修改话题为 MoveIt Servo 独立节点默认监听的话题
        // 配合您的 launch 文件，节点名为 servo_node
        twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>(
            "/servo_node/delta_twist_cmds", 10);
    }

    // 初始化函数，在对象构造完成后调用
    void initialize()
    {

        rclcpp::sleep_for(std::chrono::seconds(1));
        getCurrentPositionAsReference();

        // 5. 创建定时器，25Hz运行控制循环（40ms）
        timer_ = create_wall_timer(std::chrono::milliseconds(40), std::bind(&ForceServoControlNode::control_loop, this));

        RCLCPP_INFO(this->get_logger(), "力控伺服节点已启动 (Client Mode)，正在向 /servo_node 发布指令");
    }

private:
    rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;
    rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
    
    std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
    rclcpp::TimerBase::SharedPtr timer_;

    Eigen::Vector3d current_force_;
    Eigen::Vector3d current_position_;
    Eigen::Vector3d velocity_;
    Eigen::Matrix3d M_position_;
    Eigen::Matrix3d K_position_;
    Eigen::Matrix3d B_position_;
    Eigen::Vector3d force_torque_reference_;
    Eigen::Vector3d position_reference_;  // 存储参考位置
    double dt_;
    size_t count_;

    // 获取机械臂当前的位姿并设置为参考位置
    void getCurrentPositionAsReference()
    {
        geometry_msgs::msg::TransformStamped transform;
        try {
            // 增加超时时间，避免刚启动时报错
            transform = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero, std::chrono::seconds(1));
            position_reference_.x() = transform.transform.translation.x;
            position_reference_.y() = transform.transform.translation.y;
            position_reference_.z() = transform.transform.translation.z;
            RCLCPP_INFO(this->get_logger(), "设置参考位置为: x=%.3f, y=%.3f, z=%.3f",
                        position_reference_.x(), position_reference_.y(), position_reference_.z());
        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(this->get_logger(), "初始化时无法获取TF变换: %s", ex.what());
        }
    }

    void force_callback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
    {
        // 保持您原始的映射逻辑
        current_force_.x() = -msg->wrench.force.z;
        current_force_.y() = 0.0;
        current_force_.z() = 0.0;
    }

    void control_loop()
    {
        geometry_msgs::msg::TransformStamped transform;
        try {
            transform = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
            current_position_.x() = transform.transform.translation.x;
            current_position_.y() = transform.transform.translation.y;
            current_position_.z() = transform.transform.translation.z;
        } catch (tf2::TransformException &ex) {
            // 频率很高时使用 debug 或 throttle 避免刷屏，或者保持 WARN
            RCLCPP_WARN_THROTTLE(this->get_logger(), *this->get_clock(), 2000, "无法获取TF变换: %s", ex.what());
            return;
        }

        Eigen::Vector3d force_difference = current_force_ - force_torque_reference_;
        Eigen::Vector3d position_difference = current_position_ - position_reference_;
        
        Eigen::Vector3d kp_position_term = K_position_ * position_difference;
        Eigen::Vector3d b_velocity_term = B_position_ * velocity_;
        
        // F = Ma + Bv + Kx  =>  a = M^-1 * (F - Kx - Bv)
        Eigen::Vector3d acceleration = M_position_.inverse() * (force_difference - kp_position_term - b_velocity_term);
        velocity_ += acceleration * dt_;

        auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
        msg->header.stamp = this->now();
        msg->header.frame_id = "base_link"; // 速度是相对于 base_link 的
        
        msg->twist.linear.x = velocity_.x();
        msg->twist.linear.y = velocity_.y();
        msg->twist.linear.z = velocity_.z();
        msg->twist.angular.x = 0.0;
        msg->twist.angular.y = 0.0;
        msg->twist.angular.z = 0.0;

        twist_cmd_pub_->publish(std::move(msg));

        count_++;
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<ForceServoControlNode>();
    node->initialize();  // 在对象构造完成后调用初始化
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}